import os
import dspy
from dotenv import load_dotenv
from flask import Flask, request, jsonify, render_template
from openai import AzureOpenAI

# Default deployment to use if not specified in environment variables
DEFAULT_DEPLOYMENT = "gpt-4o-mini"

# Load environment variables
load_dotenv(".env.local")
# Fallback to .env if .env.local doesn't exist
if not os.path.exists(".env.local") and os.path.exists(".env"):
    load_dotenv(".env")

# Configure Azure OpenAI as the language model
api_key = os.getenv("AZURE_OPENAI_API_KEY")
endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
api_version = os.getenv("AZURE_OPENAI_API_VERSION", "2023-05-15")
deployment = os.getenv("AZURE_OPENAI_DEPLOYMENT", DEFAULT_DEPLOYMENT)

# Set up Azure OpenAI configuration
azure_config = {
    "api_key": api_key,
    "azure_endpoint": endpoint,
    "api_version": api_version,
    "azure_deployment": deployment
}

# Initialize the Flask application
app = Flask(__name__)

# Initialize the language model
azure_llm = None
qa_module = None
rag_model = None

# Define a simple information retrieval dataset
class Document:
    def __init__(self, text, id=None):
        self.text = text
        self.id = id

# Sample knowledge base
knowledge_base = [
    Document("DSPy is a framework for programming with foundation models. It helps developers solve complex tasks with LLMs.", "doc1"),
    Document("Azure OpenAI Service provides REST API access to OpenAI's models with the security capabilities of Azure.", "doc2"),
    Document("RAG stands for Retrieval Augmented Generation. It combines retrieval systems with LLM generation to produce more factual responses.", "doc3"),
    Document("DSPy can be integrated with Azure OpenAI to create powerful AI applications with sophisticated prompting patterns.", "doc4"),
    Document("Prompt engineering involves designing and optimizing text prompts for language models to produce desired outputs.", "doc5"),
    Document("The DSPy Teleprompter automatically optimizes prompts for specific tasks using few-shot examples.", "doc6"),
]

# Simulate a vector database
class SimpleVectorDB:
    def __init__(self, documents):
        self.documents = documents
    
    def search(self, query, k=2):
        # In a real implementation, this would use embeddings and vector search
        # For this example, we'll use a simple keyword match
        results = []
        query_terms = query.lower().split()
        
        for doc in self.documents:
            score = sum(1 for term in query_terms if term in doc.text.lower())
            if score > 0:
                results.append((doc, score))
        
        # Sort by score and return top k
        results.sort(key=lambda x: x[1], reverse=True)
        return [doc for doc, _ in results[:k]]

# Define the QA Module
class QAModule(dspy.Module):
    def __init__(self):
        super().__init__()
        self.generate_answer = dspy.ChainOfThought("question -> answer")
    
    def forward(self, question):
        return self.generate_answer(question=question)

# Define the RAG modules
class Retriever(dspy.Module):
    def __init__(self, vector_db):
        super().__init__()
        self.vector_db = vector_db
    
    def forward(self, query):
        retrieved_docs = self.vector_db.search(query)
        return dspy.Prediction(retrieved_docs=retrieved_docs)

class RAGModel(dspy.Module):
    def __init__(self, retriever):
        super().__init__()
        self.retriever = retriever
        self.generate_answer = dspy.ChainOfThought("context, question -> answer")
    
    def forward(self, question):
        # Retrieve relevant documents
        retrieval_results = self.retriever(question)
        
        # Format the context from retrieved documents
        context = "\n".join([f"Document {i+1}: {doc.text}" for i, doc in enumerate(retrieval_results.retrieved_docs)])
        
        # Generate answer using the retrieved context
        response = self.generate_answer(context=context, question=question)
        
        return dspy.Prediction(
            question=question,
            context=context,
            answer=response.answer
        )

# Initialize models
def init_models():
    global azure_llm, qa_module, rag_model
    
    # Check if credentials are available - be more explicit about checking
    # both api_key and endpoint are required
    if not api_key or not endpoint:
        print("Missing required credentials: AZURE_OPENAI_API_KEY and/or AZURE_OPENAI_ENDPOINT")
        return False
    
    try:
        # Create DSPy language model - avoid creating client directly
        azure_llm = dspy.LM(
            model=deployment,
            provider="azure",
            config={
                "api_key": api_key,
                "api_version": api_version,
                "azure_endpoint": endpoint
            }
        )
        dspy.settings.configure(lm=azure_llm)
        
        # Initialize QA Module
        qa_module = QAModule()
        
        # Initialize RAG Model
        vector_db = SimpleVectorDB(knowledge_base)
        retriever = Retriever(vector_db)
        rag_model = RAGModel(retriever)
        
        return True
    except Exception as e:
        print(f"Error initializing models: {e}")
        return False

# HTML template for the home page
@app.route('/')
def home():
    # Very simple HTML for demo purposes
    return """
    <!DOCTYPE html>
    <html>
    <head>
        <title>DSPy + Azure OpenAI Demo</title>
        <style>
            body { font-family: Arial, sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; }
            .container { margin-top: 20px; }
            .form-group { margin-bottom: 15px; }
            label { display: block; margin-bottom: 5px; }
            input[type="text"] { width: 100%; padding: 8px; box-sizing: border-box; }
            button { padding: 10px 15px; background-color: #007BFF; color: white; border: none; cursor: pointer; }
            .result { margin-top: 20px; padding: 15px; border: 1px solid #ddd; border-radius: 4px; }
            .context { margin-top: 10px; padding: 10px; background-color: #f5f5f5; border-radius: 4px; }
        </style>
    </head>
    <body>
        <h1>DSPy + Azure OpenAI Demo</h1>
        
        <div class="container">
            <h2>Ask a Question</h2>
            <div class="form-group">
                <label for="question">Question:</label>
                <input type="text" id="question" placeholder="Enter your question here...">
            </div>
            
            <div class="form-group">
                <label>Model Type:</label>
                <input type="radio" id="qa" name="model_type" value="qa" checked>
                <label for="qa">Simple QA</label>
                <input type="radio" id="rag" name="model_type" value="rag">
                <label for="rag">RAG</label>
            </div>
            
            <button onclick="askQuestion()">Submit</button>
        </div>
        
        <div id="result" class="result" style="display: none;">
            <h3>Answer:</h3>
            <p id="answer"></p>
            <div id="context-container" style="display: none;">
                <h4>Context Used:</h4>
                <div id="context" class="context"></div>
            </div>
        </div>
        
        <script>
            async function askQuestion() {
                const question = document.getElementById('question').value;
                if (!question) {
                    alert('Please enter a question');
                    return;
                }
                
                const modelType = document.querySelector('input[name="model_type"]:checked').value;
                
                try {
                    const response = await fetch('/api/ask', {
                        method: 'POST',
                        headers: {
                            'Content-Type': 'application/json'
                        },
                        body: JSON.stringify({
                            question: question,
                            model_type: modelType
                        })
                    });
                    
                    const data = await response.json();
                    
                    document.getElementById('answer').textContent = data.answer;
                    
                    // Show/hide context based on model type
                    const contextContainer = document.getElementById('context-container');
                    if (modelType === 'rag' && data.context) {
                        document.getElementById('context').textContent = data.context;
                        contextContainer.style.display = 'block';
                    } else {
                        contextContainer.style.display = 'none';
                    }
                    
                    document.getElementById('result').style.display = 'block';
                } catch (error) {
                    console.error('Error:', error);
                    alert('An error occurred. Please check the console for details.');
                }
            }
        </script>
    </body>
    </html>
    """

# Add a status route to check configuration
@app.route('/status')
def status():
    # Check environment variables
    missing_vars = []
    if not api_key:
        missing_vars.append("AZURE_OPENAI_API_KEY")
    if not endpoint:
        missing_vars.append("AZURE_OPENAI_ENDPOINT")
    
    # Prepare status information
    status_info = {
        "app_name": "DSPy + Azure OpenAI Demo",
        "config": {
            "azure_endpoint": endpoint[:10] + "..." if endpoint else "Not set",
            "api_version": api_version,
            "deployment": deployment,
            "api_key_set": api_key is not None and len(api_key) > 0
        },
        "missing_variables": missing_vars,
        "models_initialized": azure_llm is not None and qa_module is not None and rag_model is not None
    }
    
    # Return a simple HTML page showing status
    html = f"""
    <!DOCTYPE html>
    <html>
    <head>
        <title>DSPy + Azure OpenAI Demo - Status</title>
        <style>
            body {{ font-family: Arial, sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; }}
            .container {{ margin-top: 20px; }}
            .status {{ margin-top: 20px; padding: 15px; border: 1px solid #ddd; border-radius: 4px; }}
            .error {{ color: red; }}
            .success {{ color: green; }}
        </style>
    </head>
    <body>
        <h1>DSPy + Azure OpenAI Demo - Status</h1>
        
        <div class="status">
            <h2>Configuration Status</h2>
            <p><strong>Azure Endpoint:</strong> {status_info['config']['azure_endpoint']}</p>
            <p><strong>API Version:</strong> {status_info['config']['api_version']}</p>
            <p><strong>Deployment:</strong> {status_info['config']['deployment']}</p>
            <p><strong>API Key Set:</strong> {'Yes' if status_info['config']['api_key_set'] else 'No'}</p>
            
            {'<p class="error"><strong>Missing Variables:</strong> ' + ', '.join(missing_vars) + '</p>' if missing_vars else '<p class="success">All required variables are set.</p>'}
            
            <p><strong>Models Initialized:</strong> {'Yes' if status_info['models_initialized'] else 'No'}</p>
            
            {'<p>Try to initialize models now: <a href="/init">Initialize</a></p>' if not status_info['models_initialized'] else ''}
            
            <p><a href="/">Back to Home</a></p>
        </div>
    </body>
    </html>
    """
    
    return html

# Add a route to manually initialize models
@app.route('/init')
def initialize_models():
    success = init_models()
    return f"""
    <!DOCTYPE html>
    <html>
    <head>
        <title>DSPy + Azure OpenAI Demo - Initialization</title>
        <style>
            body {{ font-family: Arial, sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; }}
            .container {{ margin-top: 20px; }}
            .result {{ margin-top: 20px; padding: 15px; border: 1px solid #ddd; border-radius: 4px; }}
            .error {{ color: red; }}
            .success {{ color: green; }}
        </style>
    </head>
    <body>
        <h1>DSPy + Azure OpenAI Demo - Initialization</h1>
        
        <div class="result">
            <h2>Initialization Result</h2>
            <p class="{'success' if success else 'error'}">
                {'Models initialized successfully!' if success else 'Failed to initialize models. Check your environment variables and logs.'}
            </p>
            <p><a href="/status">Check Status</a> | <a href="/">Back to Home</a></p>
        </div>
    </body>
    </html>
    """

# API endpoint for asking questions
@app.route('/api/ask', methods=['POST'])
def ask():
    # Check if models are initialized
    if azure_llm is None:
        if not init_models():
            return jsonify({
                'error': 'Models could not be initialized. Check the /status page for more information.',
                'details': 'Missing required Azure OpenAI credentials or connection failed'
            }), 500
    
    # Get the request data
    data = request.json
    question = data.get('question')
    model_type = data.get('model_type', 'qa')
    
    if not question:
        return jsonify({'error': 'No question provided'}), 400
    
    try:
        # Process the question based on the model type
        if model_type == 'qa':
            # Use simple QA module
            response = qa_module(question)
            return jsonify({
                'question': question,
                'answer': response.answer
            })
        elif model_type == 'rag':
            # Use RAG model
            response = rag_model(question)
            return jsonify({
                'question': question,
                'answer': response.answer,
                'context': response.context
            })
        else:
            return jsonify({'error': 'Invalid model type'}), 400
    except Exception as e:
        return jsonify({'error': str(e)}), 500

# Main entry point
if __name__ == "__main__":
    # Initialize models on startup
    init_models()
    # Start the Flask app
    app.run(debug=True, host='0.0.0.0', port=5000) 